import os
import cv2
import numpy as np
import random
random.seed(20050803)
def save_img(img_list, root_dir,dir):
    if not os.path.exists(os.path.join(root_dir, dir, 'images')):
        os.makedirs(os.path.join(root_dir, dir, 'images'))
    if not os.path.exists(os.path.join(root_dir, dir, 'GT')):
        os.makedirs(os.path.join(root_dir, dir, 'GT'))
    for img_name in img_list:
        img_path = os.path.join(root_dir, 'all', 'images', img_name+'.jpg')
        GT_path = os.path.join(root_dir, 'all', 'GT', img_name+'.png')
        img = cv2.imread(img_path)
        GT = cv2.imread(GT_path, cv2.IMREAD_GRAYSCALE)
        cv2.imwrite(os.path.join(root_dir, dir, 'images', img_name+'.jpg'), img)
        cv2.imwrite(os.path.join(root_dir, dir, 'GT', img_name+'.png'), GT)

def make_dataset(root_dir, test_ratio=0.3, val_ratio=0.2):
    img_list = os.listdir(os.path.join(root_dir, 'all', 'images'))
    img_list = [img.split('.')[0] for img in img_list]
    random.shuffle(img_list)
    num_test = int(len(img_list) * test_ratio)
    num_val = int(len(img_list) * val_ratio)
    num_train = len(img_list) - num_test - num_val
    trainval_list = img_list[:num_train+num_val]
    val_list = img_list[num_train:num_train+num_val]
    test_list = img_list[num_train+num_val:]
    val_dir = 'val'
    test_dir = 'test'
    trainval_dir = 'trainval'
    save_img(trainval_list, root_dir, trainval_dir)
    save_img(val_list, root_dir, val_dir)
    save_img(test_list, root_dir, test_dir)
if __name__ == '__main__':
    root_dir = './datasets/sal/ECSSD'
    make_dataset(root_dir)
